import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import load_iris
from sklearn.svm import SVC

plt.rcParams['font.sans-serif'] = ['SimHei']     # 设置中文字体
plt.rcParams['axes.unicode_minus'] = False       # 正常显示负号

# 加载数据（两类 & 前两个特征）
data = load_iris()
X = data.data[:, :2]
y = data.target
X = X[y != 2]
y = y[y != 2]
y = np.where(y == 0, -1, 1)  # 标签改为 -1 和 +1

# 感知机（使用前面定义的）
class Perceptron:
    def __init__(self, lr=0.01, n_iters=1000):
        self.lr = lr
        self.n_iters = n_iters
        self.weights = None
        self.bias = None

    def fit(self, X, y):
        n_samples, n_features = X.shape
        self.weights = np.zeros(n_features)
        self.bias = 0
        for _ in range(self.n_iters):
            for xi, yi in zip(X, y):
                linear_output = np.dot(xi, self.weights) + self.bias
                y_pred = np.sign(linear_output)
                if yi * y_pred <= 0:
                    self.weights += self.lr * yi * xi
                    self.bias += self.lr * yi

    def predict(self, X):
        return np.sign(np.dot(X, self.weights) + self.bias)

# 训练感知机
perceptron = Perceptron(lr=0.1, n_iters=300)
perceptron.fit(X, y)
w_p, b_p = perceptron.weights, perceptron.bias

# 训练线性SVM
svm = SVC(kernel='linear', C=1e6)  # 近似硬间隔
svm.fit(X, y)
w_s = svm.coef_[0]
b_s = svm.intercept_[0]

# 准备绘图
x_vals = np.linspace(X[:, 0].min()-1, X[:, 0].max()+1, 100)

# 感知机边界
y_boundary_p = -(w_p[0] * x_vals + b_p) / w_p[1]

# SVM边界及间隔
y_boundary_s = -(w_s[0] * x_vals + b_s) / w_s[1]
margin = 1 / np.linalg.norm(w_s)
y_boundary_s_pos = -(w_s[0] * x_vals + b_s - 1) / w_s[1]
y_boundary_s_neg = -(w_s[0] * x_vals + b_s + 1) / w_s[1]

# 绘制
plt.figure(figsize=(8,6))
plt.scatter(X[:, 0], X[:, 1], c=y, cmap='coolwarm', edgecolors='k', s=50, label='样本')

# 感知机决策边界
plt.plot(x_vals, y_boundary_p, 'g--', label='感知机分界线')

# SVM决策边界及间隔
plt.plot(x_vals, y_boundary_s, 'b-', label='SVM分界线')
plt.plot(x_vals, y_boundary_s_pos, 'b:', label='SVM间隔边界')
plt.plot(x_vals, y_boundary_s_neg, 'b:')

# 高亮支持向量
plt.scatter(svm.support_vectors_[:,0], svm.support_vectors_[:,1],
            s=100, facecolors='none', edgecolors='k', label='支持向量')

plt.xlabel('Sepal length')
plt.ylabel('Sepal width')
plt.title('感知机 vs SVM 决策边界和间隔')
plt.legend()
plt.show()
